import numpy as np
from sklearn.preprocessing import RobustScaler, MinMaxScaler, StandardScaler
from tqdm import tqdm

from src.bean.vector_name_enum import VectorNameEnum
from src.bean.x_typing import PRE_HANDLE_TYPE
from src.model.model_common_util import save_vector, load_vectors
from src.util.common_util import printx, while_input, is_int_between


def pre_handle_num_v2(df, model_id, num_cols, pre_handle_type: PRE_HANDLE_TYPE):
    if pre_handle_type == PRE_HANDLE_TYPE.train:
        num = int(
            while_input(f"处理数值字段[1:StandardScaler|2:RobustScaler|3:MinMaxScaler]:", is_int_between,
                        (1, 3)))
        if num == 1:
            scaler = StandardScaler()
        elif num == 2:
            scaler = RobustScaler()
        else:
            scaler = MinMaxScaler()
        num_map = {}
        with tqdm(total=len(num_cols)) as pbs:
            for i, col in enumerate(num_cols):
                pbs.set_description(f"处理数值:{i + 1:02d}/{len(num_cols)}")
                df.loc[:, col] = scaler.fit_transform(np.array(df[col].tolist()).reshape(-1, 1).tolist())
                num_map[col] = scaler
                pbs.update()
        save_vector(num_map, VectorNameEnum.v2_num_map.value, model_id)
    else:
        num_map = load_vectors(VectorNameEnum.v2_num_map.value, model_id)
        if num_map is not None:
            for col in num_cols:
                scaler = num_map[col]
                df.loc[:, col] = scaler.transform(np.array(df[col].tolist()).reshape(-1, 1).tolist())
        else:
            printx("警告：缓存中没有数值", is_error=True)
